'''

python3 main.py 融合profiing文件 不融合profiing文件 融合json文件 输出路径

'''

import json
import pandas as pd
import numpy as np
import sys

fuse_node_perf = {}
origianl_node_perf = {}
fuse_node_to_origianl_nodes = {}
output = []
not_find_types = {}

def get_fuse_node_perf():
    df = pd.read_csv(sys.argv[1])
    for op_name in set(df[df["OP Type"].isin(["FusedAscBackend", "AscBackend"])]["Op Name"]):
        # print(df[df["Op Name"] == op_name][1:]["Task Duration(us)"].mean())
        df_chazhao = df[df["Op Name"] == op_name][1:]
        fuse_node_perf[op_name] = [df_chazhao["Task Duration(us)"].mean(), list(df_chazhao["OP Type"])[0],
        list(df_chazhao["Input Shapes"])[0], list(df_chazhao["Output Shapes"])[0]]


def get_origianl_node_perf():
    df = pd.read_csv(sys.argv[2])
    for op_name in set(df["Op Name"]):
        # print(df[df["Op Name"] == op_name][1:]["Task Duration(us)"].mean())
        df_chazhao = df[df["Op Name"] == op_name][1:]
        origianl_node_perf[op_name] = [df_chazhao["Task Duration(us)"].mean(), list(df_chazhao["OP Type"])[0]]

def read_json_file(file_path):
"""
读取JSON文件并解析其内容。
参数:
file_path -- JSON文件的路径。
返回:
解析后的数据（通常是字典或列表）。
"""
    with open(file_path, 'r', encoding='utf-8') as file:
        data = json.load(file)
    return data


def get_fuse_node_to_origianl_nodes():
    json_file_path = sys.argv[3] # 替换为你的JSON文件路径
    data = read_json_file(json_file_path)
    # print(type(data))
    # print(data.keys())
    # print(type(data["graph"]))
    # print(len(data["graph"]))
    # print(type(data["graph"][0]))
    # print(data["graph"][1].keys())
    for i in range(len(data["graph"])):
        # print(len(data["graph"][i]["op"]))
        for val in data["graph"][i]["op"]:
            if val["type"] == "AscBackend" or val["type"] == "FusedAscBackend":
                # print(val["name"])
                has_find1 = False
                has_find2 = False
                names = []
                types = []
                for attr in val["attr"]:
                    if attr["key"] == "_datadump_original_op_names":
                        has_find1 = True
                        names = attr["value"]["list"]["s"]
                    if attr["key"] == "_datadump_original_op_types":
                        has_find2 = True
                        types = attr["value"]["list"]["s"]
                if not has_find1 or not has_find2:
                    raise ZeroDivisionError("除数不能为零")
                if len(names) != len(types):
                    raise ZeroDivisionError("除数不能为零")
                fuse_node_to_origianl_nodes[val["name"]] = [names, types]

def main():
    get_fuse_node_perf()
    print("fuse_node_perf: ", len(fuse_node_perf))
    get_origianl_node_perf()
    print("origianl_node_perf: ", len(origianl_node_perf))
    get_fuse_node_to_origianl_nodes()
    print("fuse_node_to_origianl_nodes: ", len(fuse_node_to_origianl_nodes))
    for fuse_node_name, fuse_node_val in fuse_node_perf.items():
        if fuse_node_name not in fuse_node_to_origianl_nodes:
            raise ZeroDivisionError(fuse_node_name)
        origianl_nodes_time_list = []
        origianl_nodes_type = []
        origianl_nodes_name = []
        original_nodes_and_types = fuse_node_to_origianl_nodes[fuse_node_name]
        original_nodes = original_nodes_and_types[0]
        original_types = original_nodes_and_types[1]
        for node_index in range(len(original_nodes)):
            origianl_node = original_nodes[node_index]
            origianl_type = original_types[node_index]
            origianl_nodes_type.append(origianl_type)
            origianl_nodes_name.append(origianl_node)
            if origianl_node not in origianl_node_perf:
                origianl_nodes_time_list.append(0)
                if origianl_type not in not_find_types:
                    not_find_types[origianl_type] = 0
                else:
                    not_find_types[origianl_type] = not_find_types[origianl_type] + 1
            else:
                origianl_nodes_time_list.append(origianl_node_perf[origianl_node][0])
        total_origianl_nodes_time = sum(origianl_nodes_time_list)
        output_item = [fuse_node_name, fuse_node_val[0], total_origianl_nodes_time, ",".join([str(val) for val in origianl_nodes_time_list]),
        fuse_node_val[0] - total_origianl_nodes_time,
        fuse_node_val[0] - total_origianl_nodes_time*0.6, len(origianl_nodes_type),
        fuse_node_val[1],",".join(origianl_nodes_type), ",".join(origianl_nodes_name),
        fuse_node_val[2],
        fuse_node_val[3]]
        output.append(output_item)
    df_find = pd.DataFrame(np.array(output), columns=['fuse_node_name', 'fuse_node_time', 'origianl_nodes_total_time',
    'origianl_nodes_time', "gap", "expect_get",
    "nodes_num", "fuse_node_type", "origianl_nodes_type",
    "origianl_nodes_name", "Input Shapes", "Output Shapes"])
    df_find["gap"] = df_find["gap"].astype("float64")
    df_find["expect_get"] = df_find["expect_get"].astype("float64")
    df_find = df_find.sort_values(by='gap', ascending=False)
    df_find.to_csv(sys.argv[4] + 'compare_time_for_single_op.csv', index=False)
    print("not find info:")
    print(not_find_types)


if __name__ == '__main__':
    main()